﻿using Microsoft.Win32.SafeHandles;
using System;
using System.IO;
using System.Runtime.InteropServices;

namespace Liveforensics.RawProvider
{
    static internal class Disk
    {
        public enum MoveMethod : int
        {
            FileBegin = 0,
            FileCurrent,
            FileEnd
        }
        public static SafeFileHandle OpenFileHandle(string path)
        {
            SafeFileHandle handle = NativeMethods.CreateFileW(path, FileAccess.Read, FileShare.ReadWrite, IntPtr.Zero, FileMode.Open, 0, IntPtr.Zero);
            if (handle.IsInvalid)
            {
                throw Disk.GetIOExceptionForLastError();
            }
            return handle;
        }
        public static byte[] ReadSectors(SafeFileHandle handle, ulong StartSector, uint Count)
        {
            byte[] TheBuffer;
            int BytesToRead;
            long StartPoint;
            try
            {
                NativeMethods.DiskGeometry Geo = GetDiskGeometry(handle);
                BytesToRead = (int)Count * Geo.BytesPerSector;
                StartPoint = (long)StartSector * Geo.BytesPerSector;
                TheBuffer = new byte[BytesToRead];
            }
            catch
            {
                throw Disk.GetIOExceptionForLastError();
            }
            GCHandle BufHandle = GCHandle.Alloc(TheBuffer, GCHandleType.Pinned);
            
            try
            {
                
                BufHandle = GCHandle.Alloc(TheBuffer, GCHandleType.Pinned);
                IntPtr basePtr = BufHandle.AddrOfPinnedObject();
                object BytesRead = 0;
                NativeMethods.SetFilePointerEx(handle, StartPoint, null, (int)MoveMethod.FileBegin);
                bool res = NativeMethods.ReadFile(handle, basePtr, BytesToRead, BytesRead, IntPtr.Zero);
                if (BytesToRead != (int)BytesRead)
                    return null;
            }
            catch
            {
                throw Disk.GetIOExceptionForLastError();
            }
            finally
            {
                BufHandle.Free();                
            }
            return TheBuffer;
        }
        public static uint GetBytesPerSector(string path)
        {
            try
            {
                SafeFileHandle handle = OpenFileHandle(path);
                NativeMethods.DiskGeometry geo = GetDiskGeometry(handle);
                handle.Close();
                return (uint)geo.BytesPerSector;
            }
            catch { return 0; }
        }
        public static NativeMethods.DiskGeometry GetDiskGeometry(SafeFileHandle handle)
        {
            NativeMethods.DiskGeometry diskGeometry = new NativeMethods.DiskGeometry();
            int bytesRet = Marshal.SizeOf(diskGeometry);
            if (!NativeMethods.DeviceIoControl(handle, NativeMethods.EIOControlCode.DiskGetDriveGeometry, null, 0, diskGeometry, bytesRet, ref bytesRet, IntPtr.Zero))
            {
                throw Disk.GetIOExceptionForLastError();
            }
            return diskGeometry;
        }
        public static long GetDiskCapacity(SafeFileHandle diskHandle)
        {
            byte[] sizeBytes = new byte[8];
            int bytesRet = sizeBytes.Length;
            if (!NativeMethods.DeviceIoControl(diskHandle, NativeMethods.EIOControlCode.DiskGetLengthInfo, null, 0, sizeBytes, bytesRet, ref bytesRet, IntPtr.Zero))
            {
                throw Disk.GetIOExceptionForLastError();
            }

            return BitConverter.ToInt64(sizeBytes, 0);
        }
        internal static Exception GetIOExceptionForLastError()
        {
            int lastError = Marshal.GetLastWin32Error();
            int lastErrorHr = Marshal.GetHRForLastWin32Error();
            return new IOException(GetMessageForError(lastError), lastErrorHr);
        }
        public static string GetMessageForError(int code)
        {
            IntPtr buffer = new IntPtr();

            try
            {
                NativeMethods.FormatMessageW(
                    NativeMethods.FORMAT_MESSAGE_ALLOCATE_BUFFER | NativeMethods.FORMAT_MESSAGE_FROM_SYSTEM | NativeMethods.FORMAT_MESSAGE_IGNORE_INSERTS,
                    IntPtr.Zero,
                    code,
                    0,
                    ref buffer,
                    0,
                    IntPtr.Zero
                    );

                return Marshal.PtrToStringUni(buffer);
            }
            finally
            {
                if (buffer != IntPtr.Zero)
                {
                    NativeMethods.LocalFree(buffer);
                }
            }
        }
    }
}

